import tushare as ts
from pymongo import MongoClient
import json

class CatchData:
    def __init__(self):
        self.client = MongoClient()
        self.stockdb = self.client['stock']

    def save_tocsv(self, df, csvname):
        thename = csvname + ".csv"
        df.to_csv(csvname)

    def save_json(self, df, name):
        thename = name + ".json"
        df.to_json(name, orient='records')

    def save_db(self, df, tablename, drop=False):
        dbh = self.stockdb[tablename]
        if drop:
            dbh.drop()
        dbh.insert_many(json.loads(df.to_json(orient='records')))

    def get_companies(self):
        df = ts.get_stock_basics()
        csv_name = "companies"
        table_name = "companies"
        self.save_json(df, csv_name)
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_hs300(self):
        df = ts.get_hs300s()
        table_name = "hs300"
        csv_name = table_name
        self.save_json(df, csv_name)
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_sz50s(self):
        df = ts.get_sz50s()
        table_name = "sz50s"
        csv_name = table_name
        self.save_json(df, csv_name)
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_zz500s(self):
        df = ts.get_zz500s()
        table_name = "zz500s"
        csv_name = table_name
        self.save_json(df, csv_name)
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_new_stocks(self):
        df = ts.new_stocks()
        table_name = "new_stocks"
        csv_name = table_name + ".csv"
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_report_data(self, year, month):
        df = ts.get_report_data(int(year), int(month))
        thedb = self.stockdb["report_data_" + year + "_" + month]
        thedb.insert_many(json.loads(df.to_json(orient='records')))

    def get_profit_data(self, year, month):
        print("get profit data of ", year, " with ", month)
        df = ts.get_profit_data(int(year), int(month))
        thedb = self.stockdb["profit_data_" + year + "_" + month]
        thedb.insert_many(json.loads(df.to_json(orient='records')))
        
    ## how many days to get before today
    def get_topbefore(self, days, dbsavename):
        df = ts.cap_tops(days)
        df.to_csv("./" + dbsavename + ".csv")
        today_top = self.stockdb[dbsavename]
        today_top.insert_many(json.loads(df.to_json(orient='records')))

    def get_toptoday(self, date, savename):
        df = ts.top_list(date)
        df.to_csv("./" + savename + ".csv")
        toptoday = self.stockdb['toptoday']
        toptoday.insert_many(json.loads(df.to_json(orient='records')))

    def get_classified(self):
        df = ts.get_industry_classified()
        df.to_csv("./classified.csv")
        industries = self.stockdb['industry']
        industries.drop()
        industries.insert_many(json.loads(df.to_json(orient='records')))

    def get_concepts(self):
        df1 = ts.get_concept_classified()
        df1.to_csv("./concepts.csv")
        concepts = self.stockdb['concepts']
        concepts.drop()
        concepts.insert_many(json.loads(df1.to_json(orient='records')))

    def get_areas(self):
        df2 = ts.get_area_classified()
        self.save_json(df2, "areas.json")
        df2.to_csv("./areas.csv")
        areas = self.stockdb['areas']
        areas.drop()
        areas.insert_many(json.loads(df2.to_json(orient='records')))

    def set_news(self, data):
        news = self.stockdb['news']
        news.insert_many(json.loads(data))

    def get_today_data(self, date):
        data = ts.get_today_all()
        if len(data) == 0:
            print("Error get today data " + date)
        print("Get today all " + date + " of length " + str(len(data)))
        data.to_csv("./alltoday_" + date + ".csv")
        dbh = self.stockdb["alltoday_"+date]
        dbh.drop()
        dbh.insert_many(json.loads(data.to_json(orient='records')))

    def get_tick_data(self, code, date, csv_file):
        df = ts.get_tick_data(code, date)
        df.to_csv(csv_file)

    def get_k_data_e(self, code):
        ## get newest data from db
        newerdata = self.stockdb["history_" + code].find().sort("date")
        hist = None
        if newerdata.count() == 0:
            ## not exist, get all
            hist = ts.get_k_data(code)
        else:
            date = newerdata[newerdata.count() -1]["date"]
            print("Will get from date: ", newerdata[newerdata.count()-1]["date"])
            hist = ts.get_k_data(code, start=date)
        print("Save data of code ", code, " with length: ", len(hist))
        hist.to_csv("./" + code + ".csv")

    def get_h_data_newer(self, code):
        ## get newest data from db
        newerdata = self.stockdb["history_" + code].find().sort("date")
        hist = None
        if newerdata.count() == 0:
            ## not exist, get all
            hist = ts.get_h_data(code)
        else:
            date = newerdata[newerdata.count() -1]["date"]
            print("Will get from date: ", newerdata[newerdata.count()-1]["date"])
            hist = ts.get_hist_data(code, date)
        print("Save data of code ", code, " with length: ", len(hist))
        hist.to_csv("./" + code + ".csv")

    def get_fund_hodings(self, year, quarter):
        df = ts.fund_holdings(year, quarter)
        if len(df) == 0:
            print("Error get fund hodings " + str(year) + "," + str(quarter))
            return None
        print("Get fund hodings " + str(year) + "," + str(quarter) + ", length " + str(len(df)))
        table_name = "fund_holdings_" + str(year) + "_" + str(quarter)
        csv_name = table_name + ".csv"
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)

    def get_boxoffice_realtime(self):
        df = ts.realtime_boxoffice()
        if len(df) == 0:
            print("Error get boxoffice realtime")
            return None
        print("Get boxoffice realtime len [%d]", len(df))
        table_name = "boxoffice_realtime"
        csv_name = table_name + ".csv"
        self.save_tocsv(df, csv_name)
        self.save_db(df, table_name, True)
